#!/usr/bin/env python
"""
PNG encoding of a heatmap for hydroelastic.

Requires Python, PIL, and NumPy.
"""

import io
import os

import numpy as np
from PIL import Image


def main():
    h = 1
    w = 1024
    c = 3
    array = np.zeros((h, w, c), dtype=np.uint8)
    values = np.linspace(0, 1, w)

    red = np.clip((values - 0.25) * 4.0, 0.0, 1.0)
    green = np.clip((values - 0.5) * 4.0, 0.0, 1.0)
    blue = np.clip(1.0 - (values - 0.25) * 4.0, 0.0, 1.0)
    mask = values < 0.25
    blue[mask] = np.clip(values[mask] * 4.0, 0.0, 1.0)
    mask = values > 0.75
    blue[mask] = np.clip((values[mask] - 0.75) * 4.0, 0.0, 1.0)
    mask = values > 0.75

    array[0, :, 0] = red * 255.0
    array[0, :, 1] = green * 255.0
    array[0, :, 2] = blue * 255.0

    image = Image.fromarray(array)
    stream = io.BytesIO()
    image.save(stream, format='png')
    raw = stream.getvalue()

    out = "// AUTOGENERATED, DO NOT EDIT\n"
    out += "{\n"
    for byte in raw:
        out += str(int(byte)) + ", "
    out += "\n}\n"

    os.chdir(os.path.dirname(__file__))
    with open("heatmap_png.inc", "w") as f:
        f.write(out)

    print("This may need to get reformatted by ament_clang_format")


if __name__ == "__main__":
    main()
